#!/usr/bin/env python3
"""
测试DataSaver插件的split_data文件夹输出功能
"""

import tempfile
import os
import json
import sys

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
src_dir = os.path.join(current_dir, 'src')
sys.path.insert(0, src_dir)
sys.path.insert(0, current_dir)

from plugins.data_analysis.data_analysis import DataSplitter, DataSaver

def create_test_data():
    """创建测试数据"""
    return [
        {"name": "Alice", "age": 25, "department": "IT", "salary": 5000, "score": 85},
        {"name": "Bob", "age": 30, "department": "HR", "salary": 4500, "score": 78},
        {"name": "Charlie", "age": 35, "department": "IT", "salary": 6000, "score": 92},
        {"name": "Diana", "age": 28, "department": "Finance", "salary": 5500, "score": 88},
        {"name": "Eve", "age": 32, "department": "HR", "salary": 4800, "score": 82},
        {"name": "Frank", "age": 29, "department": "Finance", "salary": 5200, "score": 90},
        {"name": "Grace", "age": 27, "department": "IT", "salary": 5300, "score": 87},
        {"name": "Henry", "age": 31, "department": "HR", "salary": 4600, "score": 75}
    ]

def test_split_by_field_with_folder_output():
    """测试按字段拆分并使用文件夹输出"""
    print("=== 测试按字段拆分的文件夹输出 ===")
    
    # 创建测试数据
    test_data = create_test_data()
    
    # 创建临时文件夹
    with tempfile.TemporaryDirectory() as temp_dir:
        split_output_dir = os.path.join(temp_dir, "department_splits")
        
        # 配置拆分器 - 按部门字段拆分
        splitter_config = {
            "split_strategy": "field_value",
            "split_field": "department"
        }
        
        # 执行拆分
        splitter = DataSplitter("department_splitter", {"config": splitter_config})
        split_result = splitter.execute(test_data)
        
        print(f"拆分结果: {len(split_result.get('split_data', {}))} 个部门")
        
        # 配置保存器 - 使用文件夹输出
        saver_config = {
            "format": "json",
            "output_path": split_output_dir,
            "split_output_folder": True,  # 启用文件夹输出
            "file_prefix": "department",
            "include_metadata": True,
            "pretty_print": True
        }
        
        # 执行保存
        saver = DataSaver("folder_saver", {"config": saver_config})
        save_result = saver.execute(split_result)
        
        print(f"保存结果: {save_result}")
        
        # 验证输出
        if save_result.get("saved"):
            print(f"✓ 成功保存到文件夹: {save_result.get('file_path')}")
            print(f"✓ 保存了 {save_result.get('total_groups', 0)} 个分组")
            print(f"✓ 总共 {save_result.get('total_records', 0)} 条记录")
            
            # 检查生成的文件
            saved_files = save_result.get("saved_files", [])
            for file_info in saved_files:
                file_path = file_info["file_path"]
                if os.path.exists(file_path):
                    print(f"✓ 文件存在: {os.path.basename(file_path)} (记录数: {file_info['record_count']})")
                    
                    # 验证文件内容
                    with open(file_path, 'r', encoding='utf-8') as f:
                        content = json.load(f)
                        print(f"  - 数据记录数: {len(content.get('data', []))}")
                        print(f"  - 包含元数据: {'metadata' in content}")
                else:
                    print(f"✗ 文件不存在: {file_path}")
        else:
            print(f"✗ 保存失败: {save_result.get('error')}")

def test_split_by_condition_with_folder_output():
    """测试按条件拆分并使用文件夹输出"""
    print("\n=== 测试按条件拆分的文件夹输出 ===")
    
    # 创建测试数据
    test_data = create_test_data()
    
    # 创建临时文件夹
    with tempfile.TemporaryDirectory() as temp_dir:
        split_output_dir = os.path.join(temp_dir, "age_groups")
        
        # 配置拆分器 - 按年龄条件拆分
        splitter_config = {
            "split_strategy": "condition",
            "condition_field": "age",
            "conditions": {
                "young": "age < 30",
                "middle": "age >= 30 and age < 35",
                "senior": "age >= 35"
            }
        }
        
        # 执行拆分
        splitter = DataSplitter("age_splitter", {"config": splitter_config})
        split_result = splitter.execute(test_data)
        
        print(f"拆分结果: {len(split_result.get('split_data', {}))} 个年龄组")
        
        # 配置保存器 - 使用CSV格式和文件夹输出
        saver_config = {
            "format": "csv",
            "output_path": split_output_dir,
            "split_output_folder": True,  # 启用文件夹输出
            "file_prefix": "age_group",
            "include_metadata": False  # CSV不包含元数据
        }
        
        # 执行保存
        saver = DataSaver("folder_saver", {"config": saver_config})
        save_result = saver.execute(split_result)
        
        print(f"保存结果: {save_result}")
        
        # 验证输出
        if save_result.get("saved"):
            print(f"✓ 成功保存到文件夹: {save_result.get('file_path')}")
            
            # 检查生成的文件
            saved_files = save_result.get("saved_files", [])
            for file_info in saved_files:
                file_path = file_info["file_path"]
                if os.path.exists(file_path):
                    print(f"✓ CSV文件存在: {os.path.basename(file_path)} (记录数: {file_info['record_count']})")
                else:
                    print(f"✗ 文件不存在: {file_path}")
        else:
            print(f"✗ 保存失败: {save_result.get('error')}")

def test_split_by_ratio_with_folder_output():
    """测试按比例拆分并使用文件夹输出"""
    print("\n=== 测试按比例拆分的文件夹输出 ===")
    
    # 创建测试数据
    test_data = create_test_data()
    
    # 创建临时文件夹
    with tempfile.TemporaryDirectory() as temp_dir:
        split_output_dir = os.path.join(temp_dir, "ratio_splits")
        
        # 配置拆分器 - 按比例拆分
        splitter_config = {
            "split_strategy": "ratio",
            "ratios": [0.6, 0.4],  # 60% 和 40%
            "group_names": ["train", "test"]
        }
        
        # 执行拆分
        splitter = DataSplitter("ratio_splitter", {"config": splitter_config})
        split_result = splitter.execute(test_data)
        
        print(f"拆分结果: {len(split_result.get('split_data', {}))} 个分组")
        
        # 配置保存器 - 使用TXT格式和文件夹输出
        saver_config = {
            "format": "txt",
            "output_path": split_output_dir,
            "split_output_folder": True,  # 启用文件夹输出
            "file_prefix": "dataset",
            "pretty_print": True
        }
        
        # 执行保存
        saver = DataSaver("folder_saver", {"config": saver_config})
        save_result = saver.execute(split_result)
        
        print(f"保存结果: {save_result}")
        
        # 验证输出
        if save_result.get("saved"):
            print(f"✓ 成功保存到文件夹: {save_result.get('file_path')}")
            
            # 检查生成的文件
            saved_files = save_result.get("saved_files", [])
            for file_info in saved_files:
                file_path = file_info["file_path"]
                if os.path.exists(file_path):
                    print(f"✓ TXT文件存在: {os.path.basename(file_path)} (记录数: {file_info['record_count']})")
                    
                    # 显示文件内容预览
                    with open(file_path, 'r', encoding='utf-8') as f:
                        content = f.read()
                        lines = content.split('\n')
                        print(f"  - 内容预览 (前3行):")
                        for line in lines[:3]:
                            print(f"    {line}")
                        if len(lines) > 3:
                            print(f"    ... (共{len(lines)}行)")
                else:
                    print(f"✗ 文件不存在: {file_path}")
        else:
            print(f"✗ 保存失败: {save_result.get('error')}")

def test_without_folder_output():
    """测试不使用文件夹输出的情况（对比测试）"""
    print("\n=== 测试不使用文件夹输出（对比） ===")
    
    # 创建测试数据
    test_data = create_test_data()
    
    # 创建临时文件夹
    with tempfile.TemporaryDirectory() as temp_dir:
        output_file = os.path.join(temp_dir, "single_file.json")
        
        # 配置拆分器 - 按部门字段拆分
        splitter_config = {
            "split_strategy": "field_value",
            "split_field": "department"
        }
        
        # 执行拆分
        splitter = DataSplitter("department_splitter", {"config": splitter_config})
        split_result = splitter.execute(test_data)
        
        # 配置保存器 - 不使用文件夹输出（split_output_folder为False）
        saver_config = {
            "format": "json",
            "output_path": output_file,
            "split_output_folder": False,  # 禁用文件夹输出
            "include_metadata": True,
            "pretty_print": True
        }
        
        # 执行保存
        saver = DataSaver("single_file_saver", {"config": saver_config})
        save_result = saver.execute(split_result)
        
        print(f"保存结果: {save_result}")
        
        # 验证输出
        if save_result.get("saved"):
            print(f"✓ 成功保存到单个文件: {save_result.get('file_path')}")
            print(f"✓ 记录数: {save_result.get('record_count')}")
            
            # 验证文件内容
            file_path = save_result.get("file_path")
            if os.path.exists(file_path):
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = json.load(f)
                    print(f"✓ 文件包含元数据: {'metadata' in content}")
                    print(f"✓ 数据类型: {type(content.get('data'))}")
        else:
            print(f"✗ 保存失败: {save_result.get('error')}")

if __name__ == "__main__":
    print("开始测试DataSaver的split_data文件夹输出功能...")
    
    try:
        test_split_by_field_with_folder_output()
        test_split_by_condition_with_folder_output()
        test_split_by_ratio_with_folder_output()
        test_without_folder_output()
        
        print("\n=== 所有测试完成 ===")
        print("✓ DataSaver插件的split_data文件夹输出功能正常工作")
        
    except Exception as e:
        print(f"\n✗ 测试过程中出现错误: {e}")
        import traceback
        traceback.print_exc()